package com.atwy.graph.uf;

/**
 * @Author: 小王子火
 * @Date: 2023/2/8
 * 路径压缩，再度优化
 * 这样find就能以 O(1) 的时间找到某一节点的根节点，相应的，connected和union复杂度都下降为 O(1)。
 * 要做到这一点，非常简单，只需要在find中加一行代码
 * parent[x] = parent[parent[x]]
 * x指向父节点的父节点，实现路径压缩
 * 0->1->2将变成
 *          0->2
 *          1->2
 * 树变矮了
 *
 * 可见，调用find函数每次向树根遍历的同时，顺手将树高缩短了，
 * 最终所有树高都不会超过 3（union的时候树高可能达到 3）。
 *
 * Union-Find 算法的复杂度可以这样分析：构造函数初始化数据结构需要 O(N) 的时间和空间复杂度；
 * 连通两个节点union、判断两个节点的连通性connected、计算连通分量count所需的时间复杂度均为 O(1)。
 *
 */
public class UF_03 implements UFInterface{
    /** 统计连通分量个数 */
    private int count;
    /** 节点 x 的父节点是 parent[x] */
    private int[] parent;
    /** 新增一个数组记录每棵树包含的节点数 */
    private int[] size;

    public UF_03(int n) {
        count = n;
        parent = new int[n];
        size = new int[n];
        for (int i = 0; i < n; i++) {
            parent[i] = i;
            // 初始设为1
            size[i] = 1;
        }
    }

    @Override
    public void union(int p, int q) {
        int rootP = find(p);
        int rootQ = find(q);
        if (rootP == rootQ)return;
        // 将两颗树合到一起，将小树接到大树上去
        if(size[rootP] > size[rootQ]){
            parent[rootQ] = rootP;
            size[rootP] = size[rootP] + size[rootQ];
        }else {
            parent[rootP] = rootQ;
            size[rootQ] = size[rootP] + size[rootQ];
        }
        count--;
    }

    @Override
    public boolean connected(int p, int q) {
        int rootP = find(p);
        int rootQ = find(q);
        return rootP == rootQ;
    }

    @Override
    public int count() {
        return this.count;
    }

    /** 返回某个节点 x 的根节点 */
    private int find(int p) {
        // 根节点的父节点是自己
        while (p != parent[p]){
            // p的父节点更新成 父节点的父节点，压缩路径
            parent[p] = parent[parent[p]];

            p = parent[p];
        }
        return p;
    }
}
